package com.abc.service.impl;

import com.abc.config.SearchConfig;
import com.abc.dto.request.AddDocumentRequestDto;
import com.abc.dto.request.RerankRequestDto;
import com.abc.dto.request.SearchDocumentRequestDto;
import com.abc.dto.response.RerankResponseDto;
import com.abc.service.RerankService;
import com.abc.service.VectorStoreService;
import com.abc.util.MilvusUtil;
import com.abc.util.RRFUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.List;

/**
 * 向量存储服务实现类
 *
 * @author abc
 * @since 2024-01-01
 */
@Slf4j
@Service
public class VectorStoreServiceImpl implements VectorStoreService {

    @Autowired
    private MilvusUtil milvusUtil;

    @Autowired
    private RerankService rerankService;

    @Autowired
    private SearchConfig searchConfig;

    @Override
    public void initializeVectorStore() {
        milvusUtil.initializeVectorStore();
    }

    @Override
    public void add(AddDocumentRequestDto request) {
        List<Document> documentList = request.getDocList().stream().map(Document::new).toList();
        milvusUtil.addDocuments(documentList);
    }

    @Override
    public RerankResponseDto search(SearchDocumentRequestDto request) {
        // 执行向量相似度搜索
        List<Document> documents = performVectorSearch(request);

        // 如果没有搜索结果，直接返回空列表
        if (CollectionUtils.isEmpty(documents)) {
            log.info("向量搜索无结果，查询: {}", request.getQuery());
            return null;
        }

        // 使用 RRF 算法增强排序（多维度融合排序）
        List<Document> rrfEnhancedDocuments = searchConfig.getRrf().isEnabled()
                ? performRRFEnhancement(request.getQuery(), documents)
                : documents;

        // 使用重排序服务进行最终优化
        return performRerank(request.getQuery(), rrfEnhancedDocuments);
    }

    /**
     * 执行向量相似度搜索
     */
    private List<Document> performVectorSearch(SearchDocumentRequestDto request) {
        List<Document> documents;
        boolean hasFilter = request.getFilterExpression() != null && !request.getFilterExpression().isEmpty();
        if (hasFilter) {
            documents = milvusUtil.similaritySearch(request.getQuery(), request.getTopK(), request.getFilterExpression());
            log.info("执行过滤搜索 - 查询: {}, 过滤条件数: {}, 结果数: {}", request.getQuery(), request.getFilterExpression().size(), documents.size());
        } else {
            documents = milvusUtil.similaritySearch(request.getQuery(), request.getTopK());
            log.info("执行相似度搜索 - 查询: {}, 结果数: {}", request.getQuery(), documents.size());
        }
        return documents;
    }

    /**
     * 使用 RRF 算法增强排序
     */
    private List<Document> performRRFEnhancement(String query, List<Document> documents) {
        log.info("开始 RRF 增强排序 - 查询: {}, 输入文档数: {}", query, documents.size());

        // 使用配置化的 RRF 算法进行多维度融合排序
        List<Document> rrfEnhancedDocuments = RRFUtil.enhanceRanking(documents, query);

        log.info("RRF 增强排序完成 - 输出文档数: {} (k={})",
                rrfEnhancedDocuments.size(), searchConfig.getRrf().getK());

        return rrfEnhancedDocuments;
    }

    /**
     * 执行文档重排序
     */
    @Override
    public RerankResponseDto performRerank(String query, List<Document> documents) {
        // 提取文档文本
        List<String> documentTexts = documents.stream().map(Document::getText).filter(StringUtils::hasText).toList();

        if (CollectionUtils.isEmpty(documentTexts)) {
            log.warn("文档文本为空，无法进行重排序");
            return null;
        }
        // 构建重排序请求（使用配置化参数）
        RerankRequestDto requestDto = buildRerankRequest(query, documentTexts);
        // 执行重排序
        RerankResponseDto rerankResponse = rerankService.rerank(requestDto);
        log.info("重排序完成 - 查询: {}, 输入文档数: {}, 输出结果数: {}", query, documentTexts.size(), rerankResponse.getResults() != null ? rerankResponse.getResults().size() : 0);
        return rerankResponse;
    }

    /**
     * 构建重排序请求（使用配置化参数）
     */
    private RerankRequestDto buildRerankRequest(String query, List<String> documentTexts) {
        RerankRequestDto requestDto = new RerankRequestDto();
        requestDto.setQuery(query);
        requestDto.setDocuments(documentTexts);
        requestDto.setModel(searchConfig.getRerank().getDefaultModel());
        requestDto.setTopN(Math.min(documentTexts.size(), searchConfig.getRerank().getMaxResults()));
        requestDto.setReturnDocuments(searchConfig.getRerank().isReturnDocuments());
        return requestDto;
    }

}

